package util

import (
	"math/rand"
	"sync/atomic"
	"time"
)

const byteDate = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"

// Snowflake 结构体表示一个类似雪花算法的生成器
type Snowflake struct {
	// serialGenerator 用于生成序列号的原子变量
	serialGenerator atomic.Int64
	// dataCenter 数据中心编号
	dataCenter int
	// worker 工作节点编号
	worker int
	// reverse 是否反转生成的字符串
	reverse bool
}

// NewGoSnowflake 函数创建一个指定数据中心和工作节点编号以及反转标志的雪花算法生成器
func NewGoSnowflake(dataCenter, worker int, reverse bool) *Snowflake {
	return &Snowflake{
		serialGenerator: atomic.Int64{},
		dataCenter:      dataCenter,
		worker:          worker,
		reverse:         reverse,
	}
}

// NewGoSnowflakeRandom 函数创建一个随机数据中心和工作节点编号以及反转标志的雪花算法生成器
func NewGoSnowflakeRandom(reverse bool) *Snowflake {
	return NewGoSnowflake(int(rand.Float64()*31), int(rand.Float64()*31), reverse)
}

// NewSnowflakeDefault 函数创建一个默认的雪花算法生成器，使用随机数据中心和工作节点编号，并且反转标志为 true
func NewSnowflakeDefault() *Snowflake {
	return NewGoSnowflakeRandom(true)
}

// Next 函数生成一个唯一的 ID，类似于雪花算法的实现
func (g *Snowflake) Next() int64 {
	// 获取当前时间的毫秒数
	now := time.Now().UnixMilli()
	// 获取并自增序列号
	serial := g.serialGenerator.Add(1)
	// 通过位运算组合时间戳、数据中心编号、工作节点编号和序列号
	return (now&2199023255551)<<22 | int64(g.dataCenter)<<17 | int64(g.worker)<<12 | int64(serial&4095)
}

// NextStr 函数生成一个字符串形式的唯一 ID，可以指定前缀和后缀
func (g *Snowflake) NextStr(prefix, suffix string) string {
	next := g.Next()
	var sb []byte
	chars := []byte(byteDate)
	for next >= int64(len(chars)) {
		// 计算字符在字符集中的索引
		item := int(next % int64(len(chars)))
		// 缩小 next 的值
		next /= int64(len(chars))
		// 将对应字符添加到结果字符串中
		sb = append(sb, chars[item])
	}
	// 添加最后一个字符
	sb = append(sb, chars[int(next)])
	if g.reverse {
		// 如果需要反转字符串，则进行反转操作
		for i, j := 0, len(sb)-1; i < j; i, j = i+1, j-1 {
			sb[i], sb[j] = sb[j], sb[i]
		}
	}
	// 添加前缀和后缀，返回最终结果
	return prefix + string(sb) + suffix
}

// NextStrDefault 函数生成一个默认的字符串形式的唯一 ID，无前缀和后缀
func (g *Snowflake) NextStrDefault() string {
	return g.NextStr("", "")
}

// IDGenerate 结构体用于封装雪花算法生成器，提供生成特定格式 ID 的方法。
type IDGenerate struct {
	snowflake *Snowflake
}

// NewIDGenerate 函数创建一个 IDGenerate 实例，内部初始化一个默认的雪花算法生成器。
func NewIDGenerate() *IDGenerate {
	return &IDGenerate{
		snowflake: NewSnowflakeDefault(),
	}
}

// NextStrWithPrefix 函数根据传入的前缀生成一个字符串形式的唯一 ID，后缀为空。
func (ig *IDGenerate) NextStrWithPrefix(prefix string) string {
	return ig.snowflake.NextStr(prefix, "")
}
